/**
    Copyright (c) 2010 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "DXUT.h"
#include "DXHelper.h"
#include "MandelbrotCS.h"
#include "DataBufferDX11.h"

MandelbrotCS::MandelbrotCS() : m_pCS( NULL ), m_pUAV( NULL ), m_nWidth( 0 ), m_nHeight( 0 ), m_pConstantBuf( NULL )
{
}

MandelbrotCS::~MandelbrotCS()
{
}

/**
*/
HRESULT MandelbrotCS::OnCreateDevice( ID3D11Device* pd3dDevice )
{
    HRESULT hr = S_OK;

    // compile compute shaders
    V_RETURN( CompileShaders( pd3dDevice ) );

    // Setup constant buffers
    D3D11_BUFFER_DESC Desc;
    Desc.Usage = D3D11_USAGE_DYNAMIC;
    Desc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    Desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    Desc.MiscFlags = 0;    
    Desc.ByteWidth = sizeof( ConstantBuf );
    V_RETURN( pd3dDevice->CreateBuffer( &Desc, NULL, &m_pConstantBuf ) );

    return hr;
}

/**
*/
HRESULT MandelbrotCS::CompileShaders( ID3D11Device* pd3dDevice )
{
    HRESULT hr = S_OK;
    ID3DBlob* pBlob = NULL;

    // Compile the Mandelbrot Compute Shader
    V_RETURN( CompileShaderFromFile( L"MandelbrotCS.hlsl", "MandelbrotCS", "cs_5_0", &pBlob ) );

    // Create the Mandelbrot Compute Shader
    V_RETURN( pd3dDevice->CreateComputeShader( pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pCS ) );

    SAFE_RELEASE( pBlob );

    return hr;
}

HRESULT MandelbrotCS::OnResizedSwapChain( ID3D11Device* pd3dDevice, DataBufferDX11* pDataBuffer )
{
    HRESULT hr = S_OK;

    // destroy old uav
    SAFE_RELEASE( m_pUAV );

    // save width, height
    m_nWidth = pDataBuffer->GetWidth();
    m_nHeight = pDataBuffer->GetHeight();

    // create uav 
    V_RETURN( pd3dDevice->CreateUnorderedAccessView( pDataBuffer->GetpBufferGPU(), NULL, &m_pUAV ) );

    return hr;
}

void MandelbrotCS::OnFrameRender( ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dImmediateContext,
                                  double fTime, float fElapsedTime, void* pUserContext )
{
    HRESULT hr = S_OK;

    UpdateConstantBuf( m_nWidth, m_nHeight );

    // set constant buffer
    D3D11_MAPPED_SUBRESOURCE MappedResource;            
    V( pd3dImmediateContext->Map( m_pConstantBuf, 0, D3D11_MAP_WRITE_DISCARD, 0, &MappedResource ) );
    ConstantBuf* buf = (ConstantBuf*)MappedResource.pData;
    *buf = m_cb;
    pd3dImmediateContext->Unmap( m_pConstantBuf, 0 );
    ID3D11Buffer* ppCB[1] = { m_pConstantBuf };
    pd3dImmediateContext->CSSetConstantBuffers( 0, 1, ppCB );

    pd3dImmediateContext->CSSetShader( m_pCS, NULL, 0 );
    
    // compute shader output
    ID3D11UnorderedAccessView* aUAViews[ 1 ] = { m_pUAV };
    pd3dImmediateContext->CSSetUnorderedAccessViews( 0, 1, aUAViews, (UINT*)(&aUAViews) );

    // run compute shader
    pd3dImmediateContext->Dispatch( (m_nWidth+15)/16, (m_nHeight+15)/16, 1 );

    // clear uav
    ID3D11UnorderedAccessView* aUAViewsClear[ 1 ] = { NULL };
    pd3dImmediateContext->CSSetUnorderedAccessViews( 0, 1, aUAViewsClear, (UINT*)(&aUAViewsClear) );
}

/**
*/
void MandelbrotCS::OnDestroyDevice()
{
    SAFE_RELEASE( m_pCS );
    SAFE_RELEASE( m_pUAV );
    SAFE_RELEASE( m_pConstantBuf );
}